# 读取rosbag文件，将图像保存到本地
# 需要修改：bag文件路径、save_dir保存路径、topic名称
# 若完整读取图像，并记录时间戳： python3 readimg_from_rosbag.py --downsample 1 --save_stemp True --start_id 0

import argparse
import cv2
import os
import numpy as np

import rospy
import rosbag
from cv_bridge import CvBridge



def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', type=str,
                        default="/home/daybeha/Documents/dataset/ratslam/stlucia_2007.bag")
    parser.add_argument('--topics', nargs='+', default=['/stlucia/camera/image/compressed'])
    parser.add_argument('--save_dir', type=str,
                        default="/home/daybeha/Documents/dataset/ratslam/stlucia_2007_images")
    parser.add_argument('--save_stemp',type=bool, default=True, help='save image with timestamp')
    parser.add_argument('--downsample', type=int, default=5, help='sampling interval')
    parser.add_argument('--start_id', type=int, default=0)
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = get_args()
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    # 打开ROS bag文件
    bag = rosbag.Bag(args.bag)

    # 创建一个CvBridge对象，用于将ROS图像消息转换为OpenCV图像
    bridge = CvBridge()
    cnt = args.start_id

    timestamp_list = []
    bagname = os.path.basename(args.save_dir)
    timestamp_save_dir = os.path.join(os.path.dirname(args.save_dir), 'timestamp_' + bagname + '.txt')
    # 遍历bag文件中的每个消息
    for topic, msg, t in bag.read_messages(topics=args.topics):
        cnt += 1
        if cnt % args.downsample == 0:
            # 将ROS图像消息转换为OpenCV图像
            # cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            cv_image = bridge.compressed_imgmsg_to_cv2(msg, desired_encoding='bgr8')
            # 保存图像
            cv2.imwrite(os.path.join(args.save_dir, f"image_{cnt:05d}.png"), cv_image)
            print(f"image_{cnt:05d}.png saved! \t timestamp: {t.to_sec()}")

            if args.save_stemp:
                timestamp_list.append(t.to_sec())
    np.savetxt(timestamp_save_dir, np.array(timestamp_list))


    # 关闭ROS bag文件
    bag.close()
